import numpy as np
import matplotlib.pyplot as plt
# from scipy.signal import cont2discrete
from MPC_InputAugmentMatrix_SS_U import F2_InputAugmentMatrix_SS_U
from MPC_Matrices_PM import F4_MPC_Matrices_PM
from MPC_controller_withConstriants import F7_MPC_Controller_withConstraints
from MPC_Matrices_Constraints import F6_MPC_Matrices_Constraints
import control
## 定义系统参数
# 定义无人机质量
m = 1
# 定义重力加速度常数
g = 10

## 系统定义
# 构建系统矩阵A，n x n
A = np.array([[0, 1, 0],
        [0, 0, 1],
        [0, 0, 0,]])
# 计算A矩阵维度
n = np.size(A,axis=0)
# print(n)
# 构建输入矩阵B，n x p
B = np.array([[0], 
        [1/m],
        [0]])
p = np.size(B, axis=1)

## 系统离散
Ts = 0.1
# 连续状态空间系统
sys_c = control.ss(A, B, np.eye(n), np.zeros((n, p)))
# 连续系统转离散系统 零阶保持法离散化公式
sys_d = control.c2d(sys_c, Ts, method='zoh')
# print(sys_d)
# 提取离散系统A矩阵
A = sys_d.A
# 提取离散系统B矩阵
B = sys_d.B

## 权重设计
# 设计状态权重系数矩阵，n x n
Q = np.array( [
    [1, 0, 0],
    [0, 1, 0],
    [0, 0, 0]
])

# 设计终值权重系统矩阵，n x n
S = np.array([
    [1, 0, 0],
    [0, 1, 0],
    [0, 0, 0]
])

# 设计输入权重系数矩阵，p x p
R = np.array([
    [0.1]
])

## 系统参考值
# 系统状态参考值
xd = np.array([
    [10],
    [0],
    [-g]
])

# 构建目标转移矩阵
AD = np.eye(n)

# 计算目标输入
# 目标输入的计算: ud = B \ ((I - A) * xd)
# 相当于在 Python 中求解 B * ud = (I - A) * xd
ud = np.linalg.pinv(B) @ (np.eye(n) - A) @ xd
# print(ud)

## 初始化系统状态
x0 = np.array([
    [0],
    [0],
    [-g]
])

x = x0
# 初始化增广状态矩阵
xa = np.vstack(
    [x, xd]
)

# 初始化系统输入
u0 = np.array([[0]])
u = u0


## 系统约束定义
# 输入下限
u_low = np.array([[-3]])
# 输入上限
u_high = np.array([[2]])

# 状态下限
x_low = np.array([
    [0],
    [0],
    [-g]
])
# 状态上限
x_high = np.array([
    [10],
    [3],
    [-g]
])
# 增广状态下限
xa_low = np.vstack([
    x_low,
    [[float('-inf')],
    [float('-inf')],
    [float('-inf')]]
])
# 增广状态上限
xa_high = np.vstack([
    x_high,
    [[float('inf')],
    [float('inf')],
    [float('inf')]]
])

## 定义系统运行步数
k_steps = 100

record = dict()
# 定义x_history 零矩阵，用于存储系统状态结果，维度n x k_step
record['x1_history'] = list()
record['x2_history'] = list()
# 定义u_history 零矩阵，用于存储系统输入结果，维度p x k_step
record['u_history'] = list()

# 定义预测区间
N_P = 20

# 调用模块F2 计算系统增广矩阵 Aa, Ba, Qa, Sa, R 以及目标输入ud
# 对于非定常xd每次都要计算，应该都是在线计算的了
Aa, Ba, Qa, Sa, R, ud = F2_InputAugmentMatrix_SS_U(A, B, Q, R, S, xd)

# 调用模块F4 计算二次规划需要用到的矩阵
Phi,Gamma,Omega,Psi,F,H = F4_MPC_Matrices_PM(Aa, Ba, Qa, R, Sa, N_P)

# 调用模块F6 计算含约束二次规划需要用到的矩阵
M,Beta_bar,b = F6_MPC_Matrices_Constraints(xa_low, xa_high, u_low, u_high, N_P, Phi, Gamma)

for k in range(k_steps):
    # 调用模块F7 计算系统系统控制（输入增量）
    delta_U, delta_u = F7_MPC_Controller_withConstraints(xa,F,H,M,Beta_bar,b,p)

    # 根据输入增量计算系统输入
    u = delta_u + ud

    # 系统输入代入系统方程，计算系统响应
    x = A @ x + B @ u
    # 更新增广矩阵xa
    xa = np.vstack(
        [x, xd]
    )

    # 保存系统状态到预先定义矩阵的相应位置
    record['x1_history'].append(x[0])
    record['x2_history'].append(x[1])
    record['u_history'].append(u[0])

plt.figure
plt.plot(record['x1_history'], label='x1 position')
plt.plot(record['x2_history'], label='x2 velocity')
plt.show()

plt.figure
plt.plot(record['u_history'], label='u')
plt.show()
# 完成，结果一致